#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   video_utils.py
@Time    :   2025/01/26 10:39:44
@Author  :   Sanqi Lu
@Contact :   sanqilu376@163.com
@Desc    :   
'''

import cv2
import time
# from moviepy import VideoFileClip
import shutil
import imageio
import numpy as np
import os

def parse_video_info(video_path):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("无法打开视频文件")
        return None, None, None

    # 获取帧率
    fps = cap.get(cv2.CAP_PROP_FPS)
    # 获取视频宽度（分辨率）
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    # 获取视频高度（分辨率）
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # 根据帧率和帧数计算视频时长（秒）
    duration = frame_count / fps

    cap.release()
    return fps, (width, height), duration


def parse_video_info_moviepy(video_path):
    clip = VideoFileClip(video_path)
    fps = clip.fps
    resolution = (int(clip.w), int(clip.h))
    duration = clip.duration
    return fps, resolution, duration


def change_fps_and_resolution_but_keep_duration_imageio(input_video_path, output_video_path,
                                                        target_resolution_wh=(720, 480), num_frames=49, target_fps=8):
    """
    处理视频的函数，包括提取帧、均匀采样、调整分辨率以及保存为新视频。

    参数:
    input_video_path (str): 输入视频的路径
    output_video_path (str): 输出视频的路径
    target_resolution (tuple): 目标分辨率，格式为 (宽度, 高度)，默认为 (640, 480)
    num_frames (int): 要均匀采样的帧数，默认为49
    target_fps (int): 输出视频的帧率，默认为8
    """
    # 读取视频
    reader = imageio.get_reader(input_video_path)
    fps = reader.get_meta_data()['fps']
    duration = reader.get_meta_data()['duration']
    total_frames = int(fps * duration)

    # 均匀采样帧的索引
    frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    sampled_frames = []
    for index in frame_indices:
        frame = np.asarray(reader.get_data(index))
        # 调整分辨率
        resized_frame = cv2.resize(frame, (target_resolution_wh[0], target_resolution_wh[1]))
        sampled_frames.append(resized_frame)
    reader.close()

    # 保存为新视频
    writer = imageio.get_writer(output_video_path, fps=target_fps, macro_block_size=None)
    for frame in sampled_frames:
        writer.append_data(frame.astype(np.uint8))
    writer.close()


def extract_frames_from_video_opencv(video_path, frame_indices=[]):
    '''
    Args:
        frame_indices: 待从视频读取帧的索引列表，若为空列表或 None, 则读取全部图像帧
    '''
    try:
        # 打开视频文件
        cap = cv2.VideoCapture(video_path)
        
        if not cap.isOpened():
            raise Exception(f"无法打开视频文件: {video_path}")
        
        if frame_indices is None or len(frame_indices) == 0:
            frame_count = 0
            out_frames = []
            while True:
                try:
                    # 逐帧读取
                    ret, frame = cap.read()
                    if ret:
                        frame_count += 1
                        out_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                    else:
                        # print(f"读取帧完成，截止当前共读取帧数: {frame_count}.")
                        break

                except Exception as e:
                    print(f"处理帧时发生错误: {e}")
        else:
            # 获取视频的帧率
            fps = cap.get(cv2.CAP_PROP_FPS)
            out_frames = []
            for index in frame_indices:
                # 设置视频的当前帧位置
                cap.set(cv2.CAP_PROP_POS_FRAMES, index)
                
                # 读取帧
                ret, frame = cap.read()
                
                if ret:
                    out_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                else:
                    print(f"当前帧读取失败: {index}.")

        # 释放视频对象
        cap.release()
        return out_frames
    
    except Exception as e:
        print(f"发生错误: {e}")
    
    

def merge_imgs_to_video_opencv(imgs, output_video_path, fps):
    # 获取图像帧列表
    # 读取第一帧以获取帧的大小
    first_frame = cv2.cvtColor(imgs[0], cv2.COLOR_BGR2RGB)
    height, width, layers = first_frame.shape

    # 设置视频编码器和帧率
    fourcc = cv2.VideoWriter_fourcc(*'avc1')  # 编码器（MP4格式）
    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    # 逐帧写入视频
    for frame in imgs:
        video_writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))

    # 释放视频写入器
    video_writer.release()
    print(f"Video saved as {output_video_path}")


def read_img_of_video(input_video_path, img_path, img_idx, target_resolution_wh):
    # 读取视频
    reader = imageio.get_reader(input_video_path)
    fps = reader.get_meta_data()['fps']
    duration = reader.get_meta_data()['duration']
    total_frames = int(fps * duration)

    frame = np.asarray(reader.get_data(img_idx))
    # 调整分辨率
    if not (target_resolution_wh[0] == frame.shape[1] and target_resolution_wh[1] == frame.shape[0]):
        resized_frame = cv2.resize(frame, (target_resolution_wh[0], target_resolution_wh[1]))
        print(f"调整分辨率从：{frame.shape} 到 wh: {target_resolution_wh}")
    else:
        resized_frame = frame

    cv2.imwrite(img_path, resized_frame)
    reader.close()


def draw_pose21_onto_image_opencv(frame, hand_pose, handtype="right", decrease_color=False):
    '''

    :param frame: [h, w, 3] rgb
    :param hand_pose: [21, 2/3]
    :param handtype:
    :return:
    '''
    assert hand_pose.shape[0] == 21
    Is = [0, 1, 2, 3, 0, 5, 6, 7, 9, 10, 11, 13, 14, 15, 0, 17, 18, 19, 5, 9, 13]
    Js = [1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 14, 15, 16, 17, 18, 19, 20, 9, 13, 17]
    framebgr = np.ascontiguousarray(frame[:, :, ::-1])
    # print(framebgr.shape)
    # 绘制点坐标, color="r", s=120
    for pi in range(hand_pose.shape[0]):
        center = (int(hand_pose[pi][0]), int(hand_pose[pi][1]))
        cv2.circle(framebgr, center, radius=2, color=(0, 0, 255), thickness=2)
    # 画线：color="#000000"
    left_color = (0, 255, 0)
    right_color = (255, 0, 0)
    if decrease_color:
        left_color = (0, 200, 0)
        right_color = (200, 0, 0)
    for idx in range(len(Is)):
        start_point = (int(hand_pose[Is[idx]][0]), int(hand_pose[Is[idx]][1]))
        end_point = (int(hand_pose[Js[idx]][0]), int(hand_pose[Js[idx]][1]))
        cv2.line(framebgr, start_point, end_point, color=right_color if handtype == "right" else left_color, thickness=2)

    frame = np.ascontiguousarray(framebgr[:, :, ::-1])
    return frame


def draw_pose_onto_video(in_video_path, hand_pose, out_video_path):
    assert os.path.exists(in_video_path)
    os.makedirs(os.path.dirname(out_video_path), exist_ok=True)
    fps, resolution, duration = parse_video_info(in_video_path)
    total_frame = int(fps * duration)
    assert total_frame == hand_pose.shape[1]  # [2, t, 21, 2]
    # 读取视频
    reader = imageio.get_reader(in_video_path)
    drawn_pose_list = []
    for frame_index in range(total_frame):
        drawn_frame = np.asarray(reader.get_data(frame_index))
        for hand_idx in range(hand_pose.shape[0]):
            # 绘制pose
            drawn_frame = draw_pose21_onto_image_opencv(drawn_frame, hand_pose[hand_idx, frame_index],
                                                        handtype="right" if hand_idx == 0 else "left")

        drawn_pose_list.append(drawn_frame)

    reader.close()
    # 保存为新视频
    writer = imageio.get_writer(out_video_path, fps=fps, macro_block_size=None)
    for frame in drawn_pose_list:
        writer.append_data(frame.astype(np.uint8))
    writer.close()
    print(f"Done!")


if __name__ == "__main__":
    # fps, resolution, duration = parse_video_info("/share/home/wuqingyao_danglingwei/datas/HOI4D_release/dev_test/image.mp4")
    # print(f"{fps}, {resolution}, {duration}")
    # fps, resolution, duration = parse_video_info("/share/home/wuqingyao_danglingwei/datas/HOI4D_release/dev_test/output_video.mp4")
    # print(f"{fps}, {resolution}, {duration}")
    # change_fps_and_resolution_but_keep_duration_imageio(
    #     "/share/home/wuqingyao_danglingwei/datas/HOI4D_renamed/ZY20210800001-H1-C1-N19-S100-s02-T1-align_rgbimage.mp4",
    #     "/share/home/wuqingyao_danglingwei/datas/train_cogvideo_data/ZY20210800001-H1-C1-N19-S100-s02-T1-align_rgbimage.mp4",
    # )
    pass
